# -*- coding: utf-8 -*-
# @Date     : 2021/4/13 09:29
# @Author   : xiaochuang.wu
# @File     : manager.py
"""进程管理, 查询"""

import psutil, ujson
from psutil import Process


class MyProcess(Process):
    def __str__(self):
        return ujson.dumps(dict(
            pid=self.pid,
            name=self.name(),
            cwd=self.cwd(),
            cmd=self.cmdline(),
            status=self.status(),
            exe=self.exe(),
            user=self.username(),
            is_running=self.is_running(),
            sub_process_num=len(self.children(recursive=True)),
            parents_process_num=len(self.parents())
        ))

    @staticmethod
    def get_pid_list():
        return psutil.pids()

    @classmethod
    def get_py_process(cls):
        p_list = []
        for pid in psutil.pids():
            p = cls(pid)
            if p.name().lower() == "python":
                p_list.append(p)
        return p_list

    @classmethod
    def get_process(cls, pid):
        return cls(pid)

    @classmethod
    def process_exists(cls, pid, name):
        pid = int(pid)
        if not psutil.pid_exists(pid):
            return False
        cmd = cls(pid).cmdline()[1]
        if cmd.count(name) > 0:
            return True
        return False

    @classmethod
    def kill_process(cls, pid, name):
        """kill signal process by pid"""
        if not psutil.pid_exists(pid):
            return True
        if not cls.process_exists(pid, name):
            raise Exception(f"want kill process which is not {name}")
        psutil.Process(pid).kill()
        return True

    @classmethod
    def kill_relative_process(cls, pid, name):
        """pid is main process, will kill it's children process too."""
        if not psutil.pid_exists(pid):
            return
        if not cls.process_exists(pid, name):
            raise Exception(f"this process is not {name}")
        try:
            p = cls(pid)
            # operation sub process
            sub_process_list = p.children(recursive=True)
            for sub_p in sub_process_list:
                sub_p.kill()

            # operation parent process
            parents_process_list = p.parents()
            for pp in parents_process_list:
                if pp.name().lower() == "python" and pp.cmdline()[1].endswith(".py"):
                    pp.kill()

            # kill self
            p.kill()
        except psutil.NoSuchProcess as ex:
            print(ex)
        except psutil.AccessDenied as ex:
            print(f"access denied: {ex}")


if __name__ == '__main__':
    # print(MyProcess.get_pid_list())
    # print(MyProcess.get_process(8301))
    # MyProcess.kill_relative_process(21837, "screen.py")
    for p in MyProcess.get_py_process():
        print(p)